#!/usr/bin/env python
# -*- coding: utf-8 -*-
#
#  algorithms.py
#
#  Copyright 2012 Kevin R <KRPent@gmail.com>
#
#  This program is free software; you can redistribute it and/or modify
#  it under the terms of the GNU General Public License as published by
#  the Free Software Foundation; either version 2 of the License, or
#  (at your option) any later version.
#
#  This program is distributed in the hope that it will be useful,
#  but WITHOUT ANY WARRANTY; without even the implied warranty of
#  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
#  GNU General Public License for more details.
#
#  You should have received a copy of the GNU General Public License
#  along with this program; if not, write to the Free Software
#  Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
#  MA 02110-1301, USA.
#
#
from operator import itemgetter
import networkx as nx
from shortest_path import single_source_dijkstra

def single_source_ksp_yen(graph, node_start, max_k=2):
    distances, paths = single_source_dijkstra(graph, node_start)

    dist_dict = {}
    path_dict = {}
    for _, p in paths.items():
        if len(p) > 1:
            node_end = p[-1]

            dist = [distances[node_end]]
            path = [paths[node_end]]

            B = []

            for k in range(1, max_k):
                for i in range(0, len(path[-1]) - 1):
                    node_spur = path[-1][i]
                    path_root = path[-1][:i+1]

                    edges_removed = []
                    for curr_path in path:
                        if len(curr_path) > i and path_root == curr_path[:i+1]:
                            try:
                                if 'weight' in graph[curr_path[i]][curr_path[i+1]]:
                                    cost = graph[curr_path[i]][curr_path[i+1]]['weight']
                                else:
                                    cost = 1.0
                                graph.remove_edge(curr_path[i], curr_path[i+1])
                            except KeyError:
                                #print 'key error!'
                                continue
                            else:
                                edges_removed.append([curr_path[i], curr_path[i+1], cost])

                    cost_spur, path_spur = single_source_dijkstra(graph, node_spur, node_end, avoid_nodes=path_root[:-1])
                    #print node_spur, node_end, cost_spur, path_spur
                    if node_end in path_spur:
                        path_total = path_root[:-1] + path_spur[node_end]
                        dist_total = distances[node_spur] + cost_spur[node_end]
                        potential_k = {'cost': dist_total, 'path': path_total}

                        if not (potential_k in B):
                            B.append(potential_k)

                    for edge in edges_removed:
                        graph.add_edge(edge[0], edge[1], weight=edge[2])

                if len(B):
                    B = sorted(B, key=itemgetter('cost'))
                    dist.append(B[0]['cost'])
                    path.append(B[0]['path'])
                    #print dist
                    #print path
                    B.pop(0)
                else:
                    break

            dist_dict[(node_start, node_end)] = dist
            path_dict[(node_start, node_end)] = path
    return dist_dict, path_dict


## @package YenKSP
# Computes K-Shortest Paths using Yen's Algorithm.
#
# Yen's algorithm computes single-source K-shortest loopless paths for a graph
# with non-negative edge cost. The algorithm was published by Jin Y. Yen in 1971
# and implores any shortest path algorithm to find the best path, then proceeds
# to find K-1 deviations of the best path.

## Computes K paths from a source to a sink in the supplied graph.
#
# @param graph A digraph of class DiGraph.
# @param start The source node of the graph.
# @param sink The sink node of the graph.
# @param K The amount of paths being computed.
#
# @retval [] Array of paths, where [0] is the shortest, [1] is the next
# shortest, and so on.
#
def ksp_yen(graph, node_start, node_end, max_k=2):
    distances, paths = single_source_dijkstra(graph, node_start)

    if node_end not in distances:
        return []

    A = [{'cost': distances[node_end],
          'path': paths[node_end]}]
    B = []

    if not A[0]['path']:
        return A

    for k in range(1, max_k):
        print(k)
        for i in range(0, len(A[-1]['path']) - 1):
            node_spur = A[-1]['path'][i]
            path_root = A[-1]['path'][:i+1]
            #print node_spur
            #print path_root

            edges_removed = []
            for path_k in A:
                curr_path = path_k['path']
                if len(curr_path) > i and path_root == curr_path[:i+1]:
                    try:
                        if 'weight' in graph[curr_path[i]][curr_path[i+1]]:
                            cost = graph[curr_path[i]][curr_path[i+1]]['weight']
                        else:
                            cost = 1.0
                        graph.remove_edge(curr_path[i], curr_path[i+1])
                    except KeyError:
                        print('key error!')
                        continue
                    else:
                        edges_removed.append([curr_path[i], curr_path[i+1], cost])

            cost_spur, path_spur = single_source_dijkstra(graph, node_spur, node_end, avoid_nodes=path_root[:-1])
            #print node_spur, node_end, cost_spur, path_spur
            if node_end in path_spur:
                path_total = path_root[:-1] + path_spur[node_end]
                dist_total = distances[node_spur] + cost_spur[node_end]
                potential_k = {'cost': dist_total, 'path': path_total}

                if not (potential_k in B):
                    B.append(potential_k)

            for edge in edges_removed:
                graph.add_edge(edge[0], edge[1], weight=edge[2])

        if len(B):
            B = sorted(B, key=itemgetter('cost'))
            A.append(B[0])
            B.pop(0)
        else:
            break

    return A

if __name__ == '__main__':
    #graph = nx.DiGraph([(0, 1), (1, 2), (3, 2), (0, 3), (0, 2)])
    #a = single_source_ksp_yen(graph, 0, max_k=2)
    #print a
    from map import Map
    m = Map('map.txt')
    a = ksp_yen(m._g, '38', '52', max_k=3)
    print(a)
